import matplotlib.pyplot as plt
import numpy as np
import torch

torch.manual_seed(42)

# 1.加载股票数据，并翻转数组以便按时间顺序排列
np_loadtxt = np.loadtxt('data-02-stock_daily.csv', delimiter=',')
data = np_loadtxt[::-1]
# 2.数据进行归一化处理，以便模型训练
from sklearn.preprocessing import MinMaxScaler

data = MinMaxScaler().fit_transform(data)
# 3.设置时间窗口大小
c = 7
# 4.通过滑动窗口方式生成训练数据
x = []
y = []
for i in range(len(data) - c):
    x.append(data[i:i + c])
    y.append(data[i + c][-1])
# 5.将数据转换为torch张量
x = torch.Tensor(x)
y = torch.Tensor(y).reshape(-1, 1)
# 6.划分训练集和测试集
from sklearn.model_selection import train_test_split

train_x, test_x, train_y, test_y = train_test_split(x, y, test_size=0.2, shuffle=False, random_state=42)
# 7.打印训练集输入的形状，以便验证数据划分是否正确
print(train_x.shape)
# 8.构建线性回归模型
model = torch.nn.Linear(in_features=x.shape[2] * x.shape[3], out_features=1)
print(model.weight)
print(model.bias)
loss_fn = torch.nn.MSELoss()
op = torch.optim.Adam(model.parameters())
# 9.训练模型
model.train()
for epoch in range(1000):
    op.zero_grad()
    h = model(train_x.reshape(-1, 35))
    loss = loss_fn(h, train_y)
    # print(loss)
    loss.backward()
    op.step()
# 10.使用模型进行预测
model.eval()
with torch.no_grad():
    h = model(test_x.reshape(-1, 35))
    plt.plot(h, c='r')
    plt.plot(test_y, c='g')
    plt.show()

# 11.绘制测试集的真实值与预测值的对比图
